# -*- coding: utf-8 -*-

import sys
import json
import subprocess
import os
import shutil
import random
import string

path = sys.path[0].split('/')
path = "/".join(path[0:-1])
sys.path.append(path)

module_path = os.path.split(os.path.realpath(__file__))[0]
from mongo_db import DB
from tokenizer import QDTokenizer

def get_data_dir():
    return os.path.join(module_path, "../data")

def gen_random_file_name():
    return ''.join(random.sample(string.ascii_lowercase, 16))

def predict_old(app, db_name, collect_name):
    db_inst = DB()
    db_inst.open("127.0.0.1", "9005")
    project_list = list(db_inst.database(db_name)[collect_name].find({"unit": {"$exists": True}}))

    predict_file = os.path.join(module_path, "../data/", collect_name)
    with open(predict_file, "w+") as f:
        qd_texts = QDTokenizer.gen_qd_texts(project_list)
        f.writelines("\r\n".join(qd_texts))
    train_model = os.path.join(module_path, "../data/model.bin")

    print("model predict", predict_file)
    cmds = ["fasttext", "predict-prob", train_model, predict_file, "2"]
    p = subprocess.Popen(cmds, shell=False, stdout=subprocess.PIPE)
    predict_results = []
    while p.poll() is None:
        line = p.stdout.readline()
        line = line.strip().decode("utf-8")
        if "__label__" not in line:
            continue
        labels = line.split("__label__")[1:]
        label_creds = []
        for label in labels:
            label = label.strip()
            pos = label.rfind(" ")
            label_creds.append({'category': label[:pos], "cred": float(label[pos + 1:])})
        predict_results.append({"result": label_creds, "id": ""})

    if len(predict_results) != len(project_list):
        print("predict fail")
        return []
    if p.returncode == 0:
        print('Subprogram success')

    else:
        print('Subprogram failed')
        return []

    for index, result in enumerate(predict_results):
        result["id"] = project_list[index]["id"]

    del project_list
    os.remove(predict_file)
    return predict_results


def predict_sub_categroy(parent_predict_map, predict_results):
    
    for parent_label in parent_predict_map.keys():
        cred_info = parent_predict_map[parent_label]['cred']
        
        filename = os.path.join("/tmp", gen_random_file_name())
        with open(filename, "w+") as f:
            qd_texts = parent_predict_map[parent_label]['data']
            f.writelines("\r\n".join(qd_texts))
            del qd_texts
        train_model = os.path.join(get_data_dir(), parent_label.replace("/", "O")+"_train_model.bin")
        print("train_model ： ", train_model)
        cmds = ["fasttext", "predict-prob", train_model, filename, "2"]

        p = subprocess.Popen(cmds, shell=False, stdout=subprocess.PIPE)

        readlines = []
        while p.poll() is None:
            line = p.stdout.readline()
            line = line.strip().decode("utf-8")
            if line == '':
                continue
            readlines.append(line)
        for index in range(0, len(readlines)):
            line = readlines[index]
            if "__label__" not in line:
                predict_results[cred_info[index]['index']] = {"result": [], "id": ""}
                continue
           
            labels = line.split("__label__")[1:]
            label_creds = []
            for label in labels:
                label = label.strip()
                #print("label : ", label)
                pos = label.rfind(" ")
                label_creds.append({'category': label[:pos], "cred": float(label[pos + 1:])})
            predict_results[cred_info[index]['index']] = {"result": label_creds, "id": ""}
        
        os.remove(filename)
        
    if p.returncode == 0:
        print('Subprogram success')

    else:
        print('Subprogram failed')
        return []


def predict(app, db_name, collect_name):
    db_inst = DB()
    db_inst.open("127.0.0.1", "10025")
    project_list = list(db_inst.database(db_name)[collect_name].find({"unit": {"$exists": True}}))

    qd_texts = []
    predict_file = os.path.join("/tmp/", collect_name)
    with open(predict_file, "w+") as f:
        qd_texts = QDTokenizer.gen_qd_texts(project_list)
        f.writelines("\r\n".join(qd_texts))
    train_model = os.path.join(get_data_dir(), "fasttext_train_model.bin")

    cmds = ["fasttext", "predict-prob", train_model, predict_file, "1"]
    p = subprocess.Popen(cmds, shell=False, stdout=subprocess.PIPE)
    predict_results = [0 for i in range(0, len(project_list))]

    print("count : ", len(project_list))
    parent_predict_map = {}

    wait_handle_lines = []
    while p.poll() is None:
        line = p.stdout.readline()
        line = line.strip().decode("utf-8")
        if line == '':
            continue
        wait_handle_lines.append(line)
    
    print("index : ", len(wait_handle_lines))
    for index in range(0, len(wait_handle_lines)):
        line = wait_handle_lines[index]
        if "__label__" not in line:
            print("异常 ： ", index, line)
            predict_results[index] = {"result": [], "id": ""}
            continue
        labels = line.split("__label__")[1:]

        for label in labels:
            label = label.strip()
            pos = label.rfind(" ")
            category = label[:pos]
            cred = float(label[pos + 1:])
            if category not in parent_predict_map:
                parent_predict_map[category] = {"data": [], "cred": []}
            parent_predict_map[category]["data"].append(qd_texts[index])
            parent_predict_map[category]["cred"].append({"cred": cred, "index": index})
            break
        predict_results[index] = True
        
    predict_sub_categroy(parent_predict_map, predict_results)

    if len(predict_results) != len(project_list):
        print("predict fail")
        return []
    if p.returncode == 0:
        print('Subprogram success')

    else:
        print('Subprogram failed')
        return []
    
    for index, result in enumerate(predict_results):
        if isinstance(result, bool) is True:
            print(index, result)
            continue
        result["id"] = project_list[index]["id"]

    del project_list
    os.remove(predict_file)
    return predict_results

if __name__ == "__main__":
    predict(None, "data_analy_predict", "57cce0b98751d41da9000006")
